#include <linux/init.h>
#include <linux/module.h>
#include <linux/moduleparam.h>
#include <linux/completion.h>
#include <linux/kthread.h>
#include <linux/fs.h>

#include <linux/module.h>
#include <linux/irq.h>
#include <linux/platform_device.h>
#include <linux/dma-mapping.h>
#include <linux/cdev.h>
#include <linux/types.h>
#include <linux/mutex.h>
#include <linux/acpi.h>
#include <crypto/internal/hash.h>
#include <linux/crypto.h>
#include <crypto/hash.h>
#include <crypto/internal/skcipher.h>
#include <crypto/internal/akcipher.h>
#include <linux/ioport.h>
#include <linux/io.h>

#include "sm3_test.h"
#include "test.h"

struct sdesc {
	struct shash_desc shash;
	char ctx[];
};

static struct sdesc *init_sdesc(struct crypto_shash *alg_shash)
{
	struct sdesc *sdesc;
	int size;

	size = sizeof(struct shash_desc) + crypto_shash_descsize(alg_shash);
	sdesc = kmalloc(size, GFP_KERNEL);
	if (!sdesc)
		return ERR_PTR(-ENOMEM);
	sdesc->shash.tfm = alg_shash;
	return sdesc;
}

static int calc_hash(struct crypto_shash *alg_shash, const unsigned char *data,
		     unsigned int datalen, unsigned char *digest)
{
	struct sdesc *sdesc;
	int ret;

	sdesc = init_sdesc(alg_shash);
	if (IS_ERR(sdesc)) {
		pr_info("can't alloc sdesc\n");
		return PTR_ERR(sdesc);
	}

	if (step && datalen > 1) {
		ret = crypto_shash_init(&sdesc->shash) ?:
			      crypto_shash_update(&sdesc->shash, data, datalen - 1) ?:
			      crypto_shash_update(&sdesc->shash, data + datalen - 1, 1) ?:
			      crypto_shash_final(&sdesc->shash, digest);

	} else {
		ret = crypto_shash_digest(&sdesc->shash, data, datalen, digest);
	}

	kfree(sdesc);
	return ret;
}
static inline int do_one_ahash_op(struct ahash_request *req, int ret)
{
	struct crypto_wait *wait = req->base.data;

	return crypto_wait_req(ret, wait);
}
static int test_hash_async(const unsigned char *key, unsigned int keylen,
			   const unsigned char *data, unsigned int datalen,
			   unsigned char *digest)
{
	struct crypto_ahash *atfm = NULL;
	struct ahash_request *req = NULL;
	int ret;
	struct scatterlist sg;
	struct crypto_wait wait;
	u8 *tmp = NULL;

	atfm = crypto_alloc_ahash(alg, 0, CRYPTO_ALG_ASYNC);
	if (IS_ERR(atfm)) {
		pr_err("alg: hash: failed to allocate transform for %s: %ld\n",
		       alg, PTR_ERR(atfm));
		return PTR_ERR(atfm);
	}

	tmp = kmalloc(datalen, GFP_KERNEL);
	if (tmp == NULL) {
		ret = -ENOMEM;
		goto out;
	}
	memcpy(tmp, data, datalen);

	if (keylen)
		crypto_ahash_setkey(atfm, key, keylen);

	alg = (char *)crypto_ahash_driver_name(atfm);
	pr_err("get alg %s\n", alg);

	req = ahash_request_alloc(atfm, GFP_KERNEL);
	if (!req) {
		pr_err("alg: hash: failed to allocate request for %s\n", alg);
		ret = -ENOMEM;
		goto out;
	}

	if (step && datalen > 1) {
		// init
		crypto_init_wait(&wait);
		ahash_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					   crypto_req_done, &wait);
		ahash_request_set_crypt(req, NULL, digest, 0);
		ret = crypto_wait_req(crypto_ahash_init(req), &wait);
		if (ret) {
			pr_err("crypto_ahash_init failed ret=%d\n", ret);
			goto out;
		}

		// update
		crypto_init_wait(&wait);
		ahash_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					   crypto_req_done, &wait);

		sg_init_table(&sg, 1);
		sg_set_buf(&sg, data, datalen - 1);
		ahash_request_set_crypt(req, &sg, digest, datalen - 1);
		ret = crypto_wait_req(crypto_ahash_update(req), &wait);
		if (ret) {
			pr_err("crypto_ahash_update failed ret=%d\n", ret);
			goto out;
		}

		// finup
		crypto_init_wait(&wait);
		ahash_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					   crypto_req_done, &wait);

		sg_init_table(&sg, 1);
		sg_set_buf(&sg, data + datalen - 1, 1);
		ahash_request_set_crypt(req, &sg, digest, 1);
		ret = crypto_wait_req(crypto_ahash_finup(req), &wait);
		if (ret) {
			pr_err("crypto_ahash_finup failed ret=%d\n", ret);
		}
	} else {
		sg_init_table(&sg, 1);
		sg_set_buf(&sg, tmp, datalen);

		crypto_init_wait(&wait);
		ahash_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
					   crypto_req_done, &wait);

		ahash_request_set_crypt(req, &sg, digest, datalen);
		ret = do_one_ahash_op(req, crypto_ahash_digest(req));
		// ret = crypto_wait_req(crypto_ahash_digest(req), req->base.data);
		if (ret) {
			pr_err("crypto_ahash_digest failed ret=%d\n", ret);
		}
	}

out:
	if (tmp)
		kfree(tmp);

	if (req)
		ahash_request_free(req);

	if (atfm)
		crypto_free_ahash(atfm);

	return ret;
}

static int test_hash_sync(const unsigned char *key, unsigned int keylen,
			  const unsigned char *data, unsigned int datalen,
			  unsigned char *digest)
{
	struct crypto_shash *alg_shash;
	char *hash_alg_name = alg;
	int ret;

	alg_shash = crypto_alloc_shash(hash_alg_name, 0, 0);
	if (IS_ERR(alg_shash)) {
		pr_info("can't alloc alg %s\n", hash_alg_name);
		return PTR_ERR(alg_shash);
	}

	if (keylen) {
		ret = crypto_shash_setkey(alg_shash, key, keylen);
		if (ret) {
			goto out;
		}
	}

	alg = (char *)crypto_shash_driver_name(alg_shash);
	pr_err("get alg %s\n", alg);

	ret = calc_hash(alg_shash, data, datalen, digest);
out:
	crypto_free_shash(alg_shash);
	return ret;
}

void sm3_test(void)
{
	int test_all = 0, ret;
	int i;
	int vector_count = sizeof(sm3_tv_template) / sizeof(sm3_tv_template[0]);
	u8 result[64] = { 0 };

	if (idx == 0xffff) {
		test_all = 1;
	} else {
		if (idx >= vector_count) {
			pr_err("Invalid idx %d , not < %d\n", idx,
			       vector_count);
			return;
		}
	}

	if (test_all) {
		for (i = 0; i < vector_count; i++) {
			if (sync)
				ret = test_hash_sync(
					NULL, 0, sm3_tv_template[i].plaintext,
					sm3_tv_template[i].psize, result);
			else
				ret = test_hash_async(
					NULL, 0, sm3_tv_template[i].plaintext,
					sm3_tv_template[i].psize, result);

			if (ret)
				return;

			if (memcmp(result, sm3_tv_template[i].digest, 32)) {
				pr_err("generate hash err(step %u)[%d]!\n",
				       step, i);
				print_buf((u8 *)"plaintext",
					  (u8 *)sm3_tv_template[i].plaintext,
					  sm3_tv_template[i].psize);
				print_buf((u8 *)"result", (u8 *)result, 32);
				print_buf((u8 *)"expected digest",
					  (u8 *)sm3_tv_template[i].digest, 32);
				return;
			}
			pr_info("generate hash success (step %u sync %u)[%d]!\n",
				step, sync, i);
		}
	} else {
		i = idx;
		{
			if (sync)
				ret = test_hash_sync(
					NULL, 0, sm3_tv_template[i].plaintext,
					sm3_tv_template[i].psize, result);
			else
				ret = test_hash_async(
					NULL, 0, sm3_tv_template[i].plaintext,
					sm3_tv_template[i].psize, result);

			if (ret)
				return;

			if (memcmp(result, sm3_tv_template[i].digest, 32)) {
				pr_err("generate hash err(step %u)[%d]!\n",
				       step, i);
				print_buf((u8 *)"plaintext",
					  (u8 *)sm3_tv_template[i].plaintext,
					  sm3_tv_template[i].psize);
				print_buf((u8 *)"result", (u8 *)result, 32);
				print_buf((u8 *)"expected digest",
					  (u8 *)sm3_tv_template[i].digest, 32);
				return;
			}
			pr_info("generate hash success (step %u)[%d]!\n", step,
				i);
		}
	}

	return;
}

void hmac_sm3_test(void)
{
	int test_all = 0, ret;
	int i;
	int vector_count =
		sizeof(hmac_sm3_tv_template) / sizeof(hmac_sm3_tv_template[0]);
	u8 result[64] = { 0 };

	if (idx == 0xffff) {
		test_all = 1;
	} else {
		if (idx >= vector_count) {
			pr_err("Invalid idx %d , not < %d\n", idx,
			       vector_count);
			return;
		}
	}

	if (test_all) {
		for (i = 0; i < vector_count; i++) {
			if (sync)
				ret = test_hash_sync(
					hmac_sm3_tv_template[i].key,
					hmac_sm3_tv_template[i].ksize,
					hmac_sm3_tv_template[i].plaintext,
					hmac_sm3_tv_template[i].psize, result);
			else
				ret = test_hash_async(
					hmac_sm3_tv_template[i].key,
					hmac_sm3_tv_template[i].ksize,
					hmac_sm3_tv_template[i].plaintext,
					hmac_sm3_tv_template[i].psize, result);

			if (ret)
				return;

			if (memcmp(result, hmac_sm3_tv_template[i].digest,
				   32)) {
				pr_err("generate hmac err(step %u)[%d]!\n",
				       step, i);
				print_buf(
					(u8 *)"plaintext",
					(u8 *)hmac_sm3_tv_template[i].plaintext,
					hmac_sm3_tv_template[i].psize);
				print_buf((u8 *)"result", (u8 *)result, 32);
				print_buf((u8 *)"expected digest",
					  (u8 *)hmac_sm3_tv_template[i].digest,
					  32);
				// return;
			} else
				pr_info("generate hmac success (step %u)[%d]!\n",
					step, i);
		}
	} else {
		i = idx;
		{
			if (sync)
				ret = test_hash_sync(
					hmac_sm3_tv_template[i].key,
					hmac_sm3_tv_template[i].ksize,
					hmac_sm3_tv_template[i].plaintext,
					hmac_sm3_tv_template[i].psize, result);
			else
				ret = test_hash_async(
					hmac_sm3_tv_template[i].key,
					hmac_sm3_tv_template[i].ksize,
					hmac_sm3_tv_template[i].plaintext,
					hmac_sm3_tv_template[i].psize, result);

			if (ret)
				return;

			if (memcmp(result, hmac_sm3_tv_template[i].digest,
				   32)) {
				pr_err("generate hmac err(step %u)[%d]!\n",
				       step, i);
				print_buf(
					(u8 *)"plaintext",
					(u8 *)hmac_sm3_tv_template[i].plaintext,
					hmac_sm3_tv_template[i].psize);
				print_buf((u8 *)"result", (u8 *)result, 32);
				print_buf((u8 *)"expected digest",
					  (u8 *)hmac_sm3_tv_template[i].digest,
					  32);
				return;
			}
			pr_info("generate hmac success (step %u)[%d]!\n", step,
				i);
		}
	}
}
